import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import os


def visualize_class_samples(save_dir):
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.CIFAR10(
        root='./data', train=True, download=True, transform=transform
    )
    classes = dataset.classes
    num_classes = len(classes)
    samples_per_class = 5

    class_samples = {cls: [] for cls in classes}
    for img_tensor, label in dataset:
        cls_name = classes[label]
        if len(class_samples[cls_name]) < samples_per_class:
            img_np = img_tensor.permute(1, 2, 0).numpy()
            class_samples[cls_name].append(img_np)

        if all(len(samples) == samples_per_class for samples in class_samples.values()):
            break

    os.makedirs(save_dir, exist_ok=True)

    plt.figure(figsize=(15, 10))
    for cls_idx, (cls_name, samples) in enumerate(class_samples.items()):
        for img_idx, img in enumerate(samples):

            ax = plt.subplot(num_classes, samples_per_class, cls_idx * samples_per_class + img_idx + 1)
            ax.imshow(img)
            ax.axis('off')  # Hide axes

            if img_idx == 0:
                ax.set_title(cls_name, fontsize=12, pad=8)

    plt.suptitle("CIFAR-10 类别样本可视化", fontsize=16, y=0.98)
    save_path = os.path.join(save_dir, "class_sample_vis.png")
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()

    print(f"类别样本可视化已保存至 {save_path}")


# Run visualization
if __name__ == "__main__":
    visualize_class_samples(save_dir="../results/")
